#!/usr/bin/env python3
"""
AKShare配置和优化模块
根据官方文档 https://akshare.akfamily.xyz/ 进行配置优化
"""

import akshare as ak
import pandas as pd
import numpy as np
import logging
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class AKShareConfig:
    """AKShare配置和优化类"""
    
    def __init__(self):
        """初始化AKShare配置"""
        self.version = ak.__version__
        logger.info(f"AKShare版本: {self.version}")
        
        # 设置请求参数（根据官方文档建议）
        self.setup_requests_config()
        
    def setup_requests_config(self):
        """配置请求参数"""
        try:
            # 设置请求头，避免被反爬虫
            import requests
            requests.packages.urllib3.disable_warnings()
            
            # 设置默认请求头
            self.headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
                'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
                'Accept-Encoding': 'gzip, deflate, br',
                'Connection': 'keep-alive',
                'Upgrade-Insecure-Requests': '1',
            }
            
            logger.info("✅ 请求配置设置完成")
            
        except Exception as e:
            logger.warning(f"⚠️ 请求配置设置失败: {e}")
    
    def test_basic_functions(self):
        """测试AKShare基础功能"""
        logger.info("🧪 开始测试AKShare基础功能...")
        
        test_results = {}
        
        # 1. 测试股票基本信息
        try:
            logger.info("测试1: 获取股票基本信息")
            stock_info = ak.stock_zh_a_spot_em()
            test_results['stock_info'] = {
                'status': 'success',
                'count': len(stock_info),
                'sample': stock_info.head(3).to_dict('records') if not stock_info.empty else []
            }
            logger.info(f"✅ 股票基本信息获取成功，共{len(stock_info)}只股票")
            
        except Exception as e:
            test_results['stock_info'] = {'status': 'failed', 'error': str(e)}
            logger.error(f"❌ 股票基本信息获取失败: {e}")
        
        # 2. 测试历史数据获取
        try:
            logger.info("测试2: 获取股票历史数据")
            # 获取平安银行(000001)的历史数据
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
            
            hist_data = ak.stock_zh_a_hist(symbol="000001", period="daily", 
                                         start_date=start_date, end_date=end_date, adjust="")
            
            test_results['hist_data'] = {
                'status': 'success',
                'count': len(hist_data),
                'columns': list(hist_data.columns),
                'date_range': f"{hist_data['日期'].min()} 至 {hist_data['日期'].max()}" if not hist_data.empty else "无数据"
            }
            logger.info(f"✅ 历史数据获取成功，共{len(hist_data)}条记录")
            
        except Exception as e:
            test_results['hist_data'] = {'status': 'failed', 'error': str(e)}
            logger.error(f"❌ 历史数据获取失败: {e}")
        
        # 3. 测试实时数据
        try:
            logger.info("测试3: 获取实时行情数据")
            real_time = ak.stock_zh_a_spot_em()
            if not real_time.empty:
                # 获取前5只股票的实时数据
                sample_data = real_time.head(5)[['代码', '名称', '最新价', '涨跌幅', '成交量']]
                test_results['real_time'] = {
                    'status': 'success',
                    'count': len(real_time),
                    'sample': sample_data.to_dict('records')
                }
                logger.info(f"✅ 实时数据获取成功，共{len(real_time)}只股票")
            else:
                test_results['real_time'] = {'status': 'failed', 'error': '无实时数据'}
                
        except Exception as e:
            test_results['real_time'] = {'status': 'failed', 'error': str(e)}
            logger.error(f"❌ 实时数据获取失败: {e}")
        
        # 4. 测试指数数据
        try:
            logger.info("测试4: 获取指数数据")
            index_data = ak.stock_zh_index_spot_em()
            test_results['index_data'] = {
                'status': 'success',
                'count': len(index_data),
                'sample': index_data.head(3)[['代码', '名称', '最新价', '涨跌幅']].to_dict('records') if not index_data.empty else []
            }
            logger.info(f"✅ 指数数据获取成功，共{len(index_data)}个指数")
            
        except Exception as e:
            test_results['index_data'] = {'status': 'failed', 'error': str(e)}
            logger.error(f"❌ 指数数据获取失败: {e}")
        
        return test_results
    
    def get_optimized_stock_data(self, symbol, period="daily", days=365):
        """
        获取优化的股票数据
        :param symbol: 股票代码 (如: 000001, 600000)
        :param period: 周期 (daily, weekly, monthly)
        :param days: 获取天数
        :return: DataFrame
        """
        try:
            # 清理股票代码
            clean_symbol = symbol.replace('.SZ', '').replace('.SH', '')
            
            # 计算日期范围
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
            
            logger.info(f"📊 获取股票数据: {clean_symbol}, 周期: {period}, 日期: {start_date} - {end_date}")
            
            # 获取历史数据
            df = ak.stock_zh_a_hist(
                symbol=clean_symbol, 
                period=period,
                start_date=start_date, 
                end_date=end_date, 
                adjust="qfq"  # 前复权
            )
            
            if df.empty:
                logger.warning(f"⚠️ 未获取到股票 {symbol} 的数据")
                return pd.DataFrame()
            
            # 数据清理和标准化
            df = self._clean_stock_data(df)
            
            logger.info(f"✅ 成功获取 {symbol} 数据，共 {len(df)} 条记录")
            return df
            
        except Exception as e:
            logger.error(f"❌ 获取股票数据失败 {symbol}: {e}")
            return pd.DataFrame()
    
    def _clean_stock_data(self, df):
        """清理和标准化股票数据"""
        try:
            # 重命名列名为英文标准格式
            column_mapping = {
                '日期': 'timestamp',
                '开盘': 'open',
                '收盘': 'close', 
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '成交额': 'amount',
                '振幅': 'amplitude',
                '涨跌幅': 'pct_change',
                '涨跌额': 'change',
                '换手率': 'turnover'
            }
            
            # 应用列名映射
            df = df.rename(columns=column_mapping)
            
            # 确保timestamp列为datetime类型
            if 'timestamp' in df.columns:
                df['timestamp'] = pd.to_datetime(df['timestamp'])
            
            # 确保数值列为float类型
            numeric_columns = ['open', 'close', 'high', 'low', 'volume', 'amount']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')
            
            # 按时间排序
            if 'timestamp' in df.columns:
                df = df.sort_values('timestamp').reset_index(drop=True)
            
            # 移除空值行
            df = df.dropna(subset=['open', 'close', 'high', 'low'])
            
            return df
            
        except Exception as e:
            logger.error(f"❌ 数据清理失败: {e}")
            return df
    
    def get_stock_info(self, symbol):
        """获取股票基本信息"""
        try:
            clean_symbol = symbol.replace('.SZ', '').replace('.SH', '')
            
            # 获取股票基本信息
            stock_info = ak.stock_individual_info_em(symbol=clean_symbol)
            
            if stock_info.empty:
                return {}
            
            # 转换为字典格式
            info_dict = {}
            for _, row in stock_info.iterrows():
                info_dict[row['item']] = row['value']
            
            return info_dict
            
        except Exception as e:
            logger.error(f"❌ 获取股票信息失败 {symbol}: {e}")
            return {}
    
    def test_data_quality(self, symbol="000001", days=30):
        """测试数据质量"""
        logger.info(f"🔍 测试数据质量: {symbol}")
        
        df = self.get_optimized_stock_data(symbol, days=days)
        
        if df.empty:
            return {"status": "failed", "error": "无数据"}
        
        quality_report = {
            "status": "success",
            "total_records": len(df),
            "date_range": {
                "start": df['timestamp'].min().strftime('%Y-%m-%d'),
                "end": df['timestamp'].max().strftime('%Y-%m-%d')
            },
            "data_completeness": {
                "missing_values": df.isnull().sum().to_dict(),
                "completeness_rate": (1 - df.isnull().sum() / len(df)).to_dict()
            },
            "price_range": {
                "min_price": float(df['low'].min()),
                "max_price": float(df['high'].max()),
                "avg_price": float(df['close'].mean())
            },
            "volume_stats": {
                "avg_volume": float(df['volume'].mean()),
                "max_volume": float(df['volume'].max()),
                "min_volume": float(df['volume'].min())
            }
        }
        
        return quality_report

def main():
    """主函数 - 运行AKShare配置和测试"""
    print("🚀 AKShare配置和测试系统")
    print("=" * 60)
    
    # 初始化配置
    config = AKShareConfig()
    
    # 运行基础功能测试
    print("\n📋 基础功能测试:")
    test_results = config.test_basic_functions()
    
    for test_name, result in test_results.items():
        status = "✅" if result['status'] == 'success' else "❌"
        print(f"{status} {test_name}: {result['status']}")
        if result['status'] == 'failed':
            print(f"   错误: {result['error']}")
    
    # 数据质量测试
    print("\n🔍 数据质量测试:")
    quality_report = config.test_data_quality()
    
    if quality_report['status'] == 'success':
        print("✅ 数据质量测试通过")
        print(f"   记录数: {quality_report['total_records']}")
        print(f"   日期范围: {quality_report['date_range']['start']} 至 {quality_report['date_range']['end']}")
        print(f"   平均价格: {quality_report['price_range']['avg_price']:.2f}")
        print(f"   平均成交量: {quality_report['volume_stats']['avg_volume']:,.0f}")
    else:
        print(f"❌ 数据质量测试失败: {quality_report['error']}")
    
    print("\n" + "=" * 60)
    print("🎉 AKShare配置和测试完成！")

if __name__ == "__main__":
    main()
